
package org.movee.spring.data.hugegraph.query;

import com.baidu.hugegraph.api.gremlin.GremlinRequest;
import com.baidu.hugegraph.driver.GraphManager;
import com.baidu.hugegraph.driver.GremlinManager;
import com.baidu.hugegraph.driver.HugeClient;
import com.baidu.hugegraph.structure.graph.Edge;
import com.baidu.hugegraph.structure.graph.Vertex;
import com.baidu.hugegraph.structure.gremlin.ResultSet;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.movee.spring.data.hugegraph.convert.MappingHugeConverter;
import org.movee.spring.data.hugegraph.gremlin.BasicCrudGremlinTraversalCreator;
import org.movee.spring.data.hugegraph.gremlin.CrudGremlinTraversalCreator;
import org.movee.spring.data.hugegraph.gremlin.GremlinTraversal;
import org.movee.spring.data.hugegraph.mapping.HugeEntityType;
import org.movee.spring.data.hugegraph.mapping.HugePersistentEntity;
import org.movee.spring.data.hugegraph.schema.SchemaInitiator;
import org.springframework.lang.NonNull;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
public class HugeTemplate implements HugeOperations {

    private final HugeClient hugeClient;
    private final MappingHugeConverter mappingConverter;
    private final CrudGremlinTraversalCreator crudCreator;
    private final SchemaInitiator schemaInitiator;

    private final GremlinManager gremlinManager;
    private final GraphManager graphManager;

    public HugeTemplate(@NonNull HugeClient hugeClient,
                        @NonNull MappingHugeConverter converter,
                        @NonNull SchemaInitiator initiator) {

        this.hugeClient = hugeClient;
        this.gremlinManager = hugeClient.gremlin();
        this.graphManager = hugeClient.graph();
        this.mappingConverter = converter;
        this.schemaInitiator = initiator;
        this.crudCreator = new BasicCrudGremlinTraversalCreator(converter.getMappingContext());
    }

    @Override
    public MappingHugeConverter getConverter() {
        return this.mappingConverter;
    }

    @Override
    public HugeClient getHugeClient() {
        return hugeClient;
    }

    @Override
    public SchemaInitiator getSchemaInitiator() {
        return schemaInitiator;
    }

    private ResultSet doRemoteTraversal(final GremlinTraversal traversal) {

        log.debug("submit traversal: {}", traversal.getTraversal());
        ResultSet rs = gremlinManager.gremlin(traversal.getTraversal()).execute();
        return rs;
    }

    private ResultSet doRemoteTraversal(final GremlinTraversal traversal, final Map<String, Object> parameters) {

        log.debug("submit traversal: {}", traversal.getTraversal());
        GremlinRequest.Builder builder = gremlinManager.gremlin(traversal.getTraversal());
        parameters.forEach(builder::binding);
        ResultSet rs = builder.execute();
        return rs;
    }

    @Override
    @SuppressWarnings("unchecked")
    public List<Object> traversal(final GremlinTraversal traversal, final Class entityClass) {

        log.debug("=====traversal: {}", traversal.getTraversal());
        ResultSet results = doRemoteTraversal(traversal);
        log.debug("=======results: {}", results.size());
        return (List<Object>) mappingConverter.convertResultSetToEntity(results, entityClass);
    }

    public List<Object> traversal(final GremlinTraversal traversal,
                                  final Class entityClass,
                                  final Map<String, Object> parameters) {

        ResultSet results = doRemoteTraversal(traversal, parameters);
        return (List<Object>) mappingConverter.convertResultSetToEntity(results, entityClass);
    }

    //
    // basic crud operations start
    //

    @Override
    @SuppressWarnings("unchecked")
    public <T> T save(T entity, Class<T> entityClass) {

        Assert.isTrue(entity != null, "entity must no be null");

        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            Vertex vertex = mappingConverter.convertEntityToVertex(entity, entityClass);
            vertex = graphManager.addVertex(vertex);
            return (T) mappingConverter.convertVertexToEntity(vertex, entity, entityClass);
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            Edge edge = mappingConverter.convertEntityToEdge(entity, entityClass);
            edge = graphManager.addEdge(edge);
            return (T) mappingConverter.convertEdgeToEntity(edge, entity, entityClass);
        }

        return null;
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> List<T> saveAll(Iterable<T> entities, Class<T> entityClass) {
        Assert.isTrue(entities != null, "entities must no be null");

        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            List<Vertex> vertices = new ArrayList<>();
            entities.forEach(entity -> {
                Vertex vertex = mappingConverter.convertEntityToVertex(entity, entityClass);
                vertices.add(vertex);
            });

            List<Vertex> retVertices = graphManager.addVertices(vertices);
            return (List<T>) retVertices.stream()
                    .map(v -> mappingConverter.convertVertexToEntity(v, null, entityClass))
                    .collect(Collectors.toList());
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            List<Edge> edges = new ArrayList<>();
            entities.forEach(entity -> {
                Edge edge = mappingConverter.convertEntityToEdge(entity, entityClass);
                edges.add(edge);
            });

            List<Edge> retEdges = graphManager.addEdges(edges, false);
            return (List<T>) retEdges.stream()
                    .map(e -> mappingConverter.convertEdgeToEntity(e, null, entityClass))
                    .collect(Collectors.toList());
        }

        return new ArrayList<>();
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> void justSaveAll(Iterable<T> entities, Class<T> entityClass) {

        Assert.isTrue(entities != null, "entities must no be null");
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            List<Vertex> vertices = new ArrayList<>();
            entities.forEach(entity -> {
                Vertex vertex = mappingConverter.convertEntityToVertex(entity, entityClass);
                vertices.add(vertex);
            });

            graphManager.addVertices(vertices);
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            List<Edge> edges = new ArrayList<>();
            entities.forEach(entity -> {
                Edge edge = mappingConverter.convertEntityToEdge(entity, entityClass);
                edges.add(edge);
            });

            graphManager.addEdges(edges, false);
        }
    }

    @Override
    public <T> T findById(Object id, Class<T> entityClass) {
        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(id));

        GremlinTraversal traversal = null;
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            traversal = GremlinTraversal.createDirectly("g.V(ids)");
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            traversal = GremlinTraversal.createDirectly("g.E(ids)");
        }

        if (traversal != null) {
            List<Object> entities = traversal(traversal, entityClass, bindings);
            return CollectionUtils.isEmpty(entities) ? null : (T) entities.get(0);
        }
        return null;
    }

    @Override
    public <T> List<T> findAllById(Iterable<Object> ids, Class<T> entityClass) {
        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(ids));

        GremlinTraversal traversal = null;
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            traversal = GremlinTraversal.createDirectly("g.V(ids)");
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            traversal = GremlinTraversal.createDirectly("g.E(ids)");
        }

        if (traversal != null) {
            List<Object> entities = traversal(traversal, entityClass, bindings);
            return (List<T>) entities;
        }
        return new ArrayList<>();
    }

    @Override
    public <T> boolean existsById(Object id, Class<T> entityClass) {
        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(id));

        GremlinTraversal traversal = null;
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        if (persistentEntity.getEntityType() == HugeEntityType.VERTEX) {
            traversal = GremlinTraversal.createDirectly("g.V(ids).id()");
        } else if (persistentEntity.getEntityType() == HugeEntityType.EDGE) {
            traversal = GremlinTraversal.createDirectly("g.E(ids).id()");
        }

        if (traversal != null) {
            List<Object> entities = traversal(traversal, entityClass, bindings);
            return !CollectionUtils.isEmpty(entities);
        }
        return false;
    }

    private <T> void deleteById(HugeEntityType type, Class<T> entityClass, Map<String, Object> bindings) {
        GremlinTraversal traversal = null;
        if (type == HugeEntityType.VERTEX) {
            traversal = GremlinTraversal.createDirectly("g.V(ids).drop()");
        } else if (type == HugeEntityType.EDGE) {
            traversal = GremlinTraversal.createDirectly("g.E(ids).drop()");
        }

        if (traversal != null) {
            traversal(traversal, entityClass, bindings);
        }
    }

    @Override
    public <T> void deleteById(Object id, Class<T> entityClass) {
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);

        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(id));

        deleteById(persistentEntity.getEntityType(), entityClass, bindings);
    }

    @Override
    public <T> void delete(T entity, Class<T> entityClass) {
        Assert.isTrue(entity != null, "entity must no be null");
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        Object id = persistentEntity.getId(entity);
        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(id));

        deleteById(persistentEntity.getEntityType(), entityClass, bindings);
    }

    @Override
    public <T> void deleteAllById(Iterable<Object> ids, Class<T> entityClass) {
        HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);

        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", Lists.newArrayList(ids));

        deleteById(persistentEntity.getEntityType(), entityClass, bindings);
    }

    @Override
    public <T> void deleteAll(Iterable<T> entities, Class<T> entityClass) {
        Assert.isTrue(entities != null, "entities must no be null");
        final HugePersistentEntity<?> persistentEntity = mappingConverter.getPersistentEntity(entityClass);
        final List<Object> ids = new ArrayList<>();
        entities.forEach(entity -> {
            Object id = persistentEntity.getId(entity);
            ids.add(id);
        });

        Map<String, Object> bindings = new HashMap<>();
        bindings.put("ids", ids);

        deleteById(persistentEntity.getEntityType(), entityClass, bindings);
    }


    //
    // basic crud operations end
    //


}

